#include <jni.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <android/log.h>

//ref:https://github.com/mackron/dr_libs/blob/master/dr_wav.h
#define DR_WAV_IMPLEMENTATION

#include "timing.h"
#include "dr_wav.h"
#include "noise_suppression.h"
#include "resampler.h"

#ifndef nullptr
#define nullptr 0
#endif

#ifndef MIN
#define MIN(A, B)        ((A) < (B) ? (A) : (B))
#endif

#define LOGV(...)   __android_log_print((int)ANDROID_LOG_INFO, "WEBTRC_NS", __VA_ARGS__)

//写wav文件
static void wavWrite_int16(const char *filename, int16_t *buffer, size_t sampleRate, size_t totalSampleCount) {
    drwav_data_format format = {};
    format.container = drwav_container_riff;     // <-- drwav_container_riff = normal WAV files, drwav_container_w64 = Sony Wave64.
    format.format = DR_WAVE_FORMAT_PCM;          // <-- Any of the DR_WAVE_FORMAT_* codes.
    format.channels = 1;
    format.sampleRate = (drwav_uint32) sampleRate;
    format.bitsPerSample = 16;
    drwav *pWav = drwav_open_file_write(filename, &format);
    if (pWav) {
        drwav_uint64 samplesWritten = drwav_write(pWav, totalSampleCount, buffer);
        drwav_uninit(pWav);
        if (samplesWritten != totalSampleCount) {
            fprintf(stderr, "ERROR\n");
            exit(1);
        }
    }
}

//读取wav文件
static int16_t *wavRead_int16(const char *filename, uint32_t *sampleRate, uint64_t *totalSampleCount) {
    unsigned int channels;
    int16_t *buffer = drwav_open_and_read_file_s16(filename, &channels, sampleRate, totalSampleCount);
    if (buffer == nullptr) {
        printf("ERROR.");
    }
    return buffer;
}

enum nsLevel {
    kLow = 0,
    kModerate,
    kHigh,
    kVeryHigh
};

int16_t *resampler(int16_t *data_in, size_t totalSampleCount, size_t in_sample_rate, size_t out_sample_rate) {
    if (data_in == nullptr)
        return nullptr;
    if (in_sample_rate == 0) return nullptr;
    size_t lengthIn = in_sample_rate / 100;
    size_t maxLen = out_sample_rate / 100;
    const int channels = 1;
    Resampler rs;
    if (rs.Reset(in_sample_rate, out_sample_rate, channels) == -1)
        return nullptr;
    size_t outLen = (size_t) (totalSampleCount * out_sample_rate / in_sample_rate);
    int16_t *data_out = (int16_t *) malloc(outLen * sizeof(int16_t));
    if (data_out == nullptr) return nullptr;
    size_t nCount = (totalSampleCount / lengthIn);
    size_t nLast = totalSampleCount - (lengthIn * nCount);
    int16_t *samplesIn = data_in;
    int16_t *samplesOut = data_out;

    outLen = 0;
    for (int i = 0; i < nCount; i++) {
        rs.Push(samplesIn, lengthIn, samplesOut, maxLen, outLen);
        samplesIn += lengthIn;
        samplesOut += outLen;
    }
    if (nLast != 0) {
        const int max_samples = 1920;
        int16_t samplePatchIn[max_samples] = {0};
        int16_t samplePatchOut[max_samples] = {0};
        memcpy(samplePatchIn, samplesIn, nLast * sizeof(int16_t));
        rs.Push(samplePatchIn, lengthIn, samplePatchOut, maxLen, outLen);
        memcpy(samplesOut, samplePatchOut, (nLast * out_sample_rate / in_sample_rate) * sizeof(int16_t));
    }
    return data_out;
}

// 降噪
static int nsProcessWithBuffer(int16_t *buffer, uint32_t sampleRate, int samplesCount, enum nsLevel level) {
    if (buffer == nullptr) return -1;
    if (samplesCount == 0) return -1;
    size_t samples = MIN(160, sampleRate / 100);
    if (samples == 0) return -1;
    uint32_t num_bands = 1;
    int16_t *input = buffer;
    size_t nTotal = (samplesCount / samples);
    NsHandle *nsHandle = WebRtcNs_Create();
    int status = WebRtcNs_Init(nsHandle, sampleRate);
    if (status != 0) {
        printf("WebRtcNs_Init fail\n");
        return -1;
    }
    status = WebRtcNs_set_policy(nsHandle, 2);
    if (status != 0) {
        printf("WebRtcNs_set_policy fail\n");
        return -1;
    }
    for (int i = 0; i < nTotal; i++) {
        int16_t *nsIn[1] = {input};   //ns input[band][data]
        int16_t *nsOut[1] = {input};  //ns output[band][data]
        WebRtcNs_Analyze(nsHandle, nsIn[0]);
        WebRtcNs_Process(nsHandle, (const int16_t *const *) nsIn, num_bands, nsOut);
        input += samples;
    }
    WebRtcNs_Free(nsHandle);

    return 1;
}

int nsProcess(int16_t *buffer, uint32_t sampleRate, uint64_t samplesCount, uint32_t channels, enum nsLevel level) {
    if (buffer == nullptr) return -1;
    if (samplesCount == 0) return -1;
    size_t samples = MIN(160, sampleRate / 100);
    if (samples == 0) return -1;
    uint32_t num_bands = 1;
    int16_t *input = buffer;
    size_t frames = (samplesCount / (samples * channels));
    int16_t *frameBuffer = (int16_t *) malloc(sizeof(*frameBuffer) * channels * samples);
    NsHandle **NsHandles = (NsHandle **) malloc(channels * sizeof(NsHandle *));
    if (NsHandles == NULL || frameBuffer == NULL) {
        if (NsHandles)
            free(NsHandles);
        if (frameBuffer)
            free(frameBuffer);
        fprintf(stderr, "malloc error.\n");
        return -1;
    }
    for (int i = 0; i < channels; i++) {
        NsHandles[i] = WebRtcNs_Create();
        if (NsHandles[i] != NULL) {
            int status = WebRtcNs_Init(NsHandles[i], sampleRate);
            if (status != 0) {
                fprintf(stderr, "WebRtcNs_Init fail\n");
                WebRtcNs_Free(NsHandles[i]);
                NsHandles[i] = NULL;
            } else {
                status = WebRtcNs_set_policy(NsHandles[i], level);
                if (status != 0) {
                    fprintf(stderr, "WebRtcNs_set_policy fail\n");
                    WebRtcNs_Free(NsHandles[i]);
                    NsHandles[i] = NULL;
                }
            }
        }
        if (NsHandles[i] == NULL) {
            for (int x = 0; x < i; x++) {
                if (NsHandles[x]) {
                    WebRtcNs_Free(NsHandles[x]);
                }
            }
            free(NsHandles);
            free(frameBuffer);
            return -1;
        }
    }
    for (int i = 0; i < frames; i++) {
        for (int c = 0; c < channels; c++) {
            for (int k = 0; k < samples; k++)
                frameBuffer[k] = input[k * channels + c];

            int16_t *nsIn[1] = {frameBuffer};   //ns input[band][data]
            int16_t *nsOut[1] = {frameBuffer};  //ns output[band][data]
            WebRtcNs_Analyze(NsHandles[c], nsIn[0]);
            WebRtcNs_Process(NsHandles[c], (const int16_t *const *) nsIn, num_bands, nsOut);
            for (int k = 0; k < samples; k++)
                input[k * channels + c] = frameBuffer[k];
        }
        input += samples * channels;
    }

    for (int i = 0; i < channels; i++) {
        if (NsHandles[i]) {
            WebRtcNs_Free(NsHandles[i]);
        }
    }
    free(NsHandles);
    free(frameBuffer);
    return 1;
}

// 降噪
static void nsProcessWithFile(const char *in_file, const char *out_file, enum nsLevel level) {
    //音频采样率
    uint32_t sampleRate = 0;
    //总音频采样数
    uint64_t inSampleCount = 0;
    int16_t *inBuffer = wavRead_int16(in_file, &sampleRate, &inSampleCount);

    //如果加载成功
    if (inBuffer != nullptr) {
        double startTime = now();
        nsProcessWithBuffer(inBuffer, sampleRate, inSampleCount, level);
        double time_interval = calcElapsed(startTime, now());
        printf("time interval: %d ms\n ", (int) (time_interval * 1000));
        wavWrite_int16(out_file, inBuffer, sampleRate, inSampleCount);
        free(inBuffer);
    }
}

extern "C"
{
JNIEXPORT void
Java_com_webrtc_ns_WebRTCNoiseSuppression_process(JNIEnv *env, jobject thiz, jstring jinputFile,
                                                  jstring joutputFile, jint level) {

    const char *inputFile = env->GetStringUTFChars(jinputFile, 0);
    const char *outputFile = env->GetStringUTFChars(joutputFile, 0);

    nsProcessWithFile(inputFile, outputFile, (enum nsLevel) level);
}

jshortArray convertShortArrayToJNI(JNIEnv *env, int16_t *arr, size_t count) {
    jshortArray retval = env->NewShortArray(count);
    env->SetShortArrayRegion(retval, 0, count, arr);
    return retval;
}


// 音频采样率 sampleRate 总音频采样数 inSampleCount
// JNIEXPORT Java_com_webrtc_ns_WebRTCNoiseSuppression_processpcm
JNIEXPORT jshortArray Java_com_webrtc_ns_WebRTCNoiseSuppression_processpcm(
        JNIEnv *env, jobject,
        jshortArray jInput,
        jint sampleRate,
        jint inSampleCount,
        jint level) {
    // 采样率从44.1到16k
    jshort *pshorts = nullptr;
    int out_sample_rate = 16000;
    int16_t *outBuffer = nullptr;
    size_t outSampleCount = 0;
    pshorts = env->GetShortArrayElements(jInput, 0);
    int short_len = env->GetArrayLength(jInput);
    int byte_len = short_len * 2;
    int16_t *pcm_use = (int16_t *) malloc(byte_len);
    memset(pcm_use, 0, byte_len);
    memcpy(pcm_use, pshorts, byte_len);

    // 采样率处理
    if (pcm_use != nullptr) {
        outBuffer = resampler(pcm_use, (size_t) inSampleCount, sampleRate, out_sample_rate);
        if (outBuffer != nullptr) {
            outSampleCount = (size_t) (inSampleCount * (out_sample_rate * 1.0f / sampleRate));

        }
    }

    nsProcessWithBuffer(outBuffer, 16000, outSampleCount, (nsLevel) level);
    jshortArray jOutput = convertShortArrayToJNI(env, pcm_use, short_len);
    env->ReleaseShortArrayElements(jInput, pshorts, 0);
    free(pcm_use);
    if (outBuffer) {
        free(outBuffer);
        outBuffer = nullptr;
    }
    return jOutput;
}

//static JNINativeMethod gMethods[] ={
//        {"processpcm", "([SIII)[S", (void*)processpcm},
//};

//
#define JNIREG_CLASS "com/webrtc/ns/WebRTCNoiseSuppression"//指定要注册的类

/*
* Register several native methods for one class.
*/
static int registerNativeMethods(JNIEnv* env, const char* className,
                                 JNINativeMethod* gMethods, int numMethods)
{
    jclass clazz;
    clazz = env->FindClass(className);
    if (clazz == NULL) {
        return JNI_FALSE;
    }
    if (env->RegisterNatives(clazz, gMethods, numMethods) < 0) {
        return JNI_FALSE;
    }

    return JNI_TRUE;
}


/*
* Register native methods for all classes we know about.
*/
//static int registerNatives(JNIEnv* env)
//{
//    if (!registerNativeMethods(env, JNIREG_CLASS, gMethods,
//                               sizeof(gMethods) / sizeof(gMethods[0])))
//        return JNI_FALSE;
//
//    return JNI_TRUE;
//}

/*
* Set some test stuff up.
*
* Returns the JNI version on success, -1 on failure.
*/
//JNIEXPORT jint JNICALL JNI_OnLoad(JavaVM* vm, void* reserved)
//{
//    JNIEnv* env = NULL;
//    jint result = -1;
//
//    if (vm->GetEnv((void**) &env, JNI_VERSION_1_4) != JNI_OK) {
//        return -1;
//    }
//    assert(env != NULL);
//
//    if (!registerNatives(env)) {//注册
//        return -1;
//    }
//    /* success -- return valid version number */
//    result = JNI_VERSION_1_4;
//
//    return result;
//}

}